#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PostgreSQL pgbench性能测试模块
使用PostgreSQL官方的pgbench工具进行数据库性能测试
"""

import os
import time
import subprocess
import re
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass

from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.progress import Progress, TaskID
from rich.box import ROUNDED

from ..core.utils import logger, run_command, detect_postgresql_version
from ..core.ui_utils import console, print_info, print_error, print_warning, print_success, get_input, confirm, ProgressManager

@dataclass
class PgbenchTestResult:
    """pgbench测试结果数据类"""
    test_type: str
    scale_factor: int
    clients: int
    threads: int
    duration: int
    tps_including_connections: float
    tps_excluding_connections: float
    latency_stddev: float
    latency_avg: float
    success: bool
    error_message: Optional[str] = None

class PgbenchTester:
    """pgbench性能测试类"""
    
    def __init__(self):
        """初始化pgbench测试器"""
        self.pg_version = detect_postgresql_version()
        self.test_db = "pgbench_test"
        self.conn_params = {}
        self.test_results = []
        
    def check_pgbench_availability(self) -> bool:
        """检查pgbench是否可用"""
        try:
            # 检查pgbench命令是否存在
            result = run_command(["pgbench", "--version"], capture_output=True)
            if result.returncode == 0:
                version_info = result.stdout.strip()
                print_info(f"检测到pgbench: {version_info}")
                return True
            else:
                print_error("未找到pgbench命令")
                return False
        except Exception as e:
            print_error(f"检查pgbench可用性失败: {e}")
            return False
    
    def get_connection_params(self) -> Dict[str, Any]:
        """获取数据库连接参数"""
        print_info("配置数据库连接...")
        
        host = get_input("请输入数据库主机", "localhost")
        port = get_input("请输入数据库端口", "5432")
        dbname = get_input("请输入数据库名", "postgres")
        user = get_input("请输入用户名", "postgres")
        password = get_input("请输入密码(留空则无密码)")
        
        try:
            port = int(port)
        except ValueError:
            print_error("无效的端口号，使用默认值5432")
            port = 5432
            
        self.conn_params = {
            'host': host,
            'port': port,
            'dbname': dbname,
            'user': user,
            'password': password
        }
        
        return self.conn_params
    
    def build_connection_string(self, dbname: Optional[str] = None) -> str:
        """构建PostgreSQL连接字符串"""
        params = self.conn_params.copy()
        if dbname:
            params['dbname'] = dbname
            
        conn_parts = []
        for key, value in params.items():
            if value:  # 只添加非空值
                if key == 'password':
                    conn_parts.append(f"PGPASSWORD={value}")
                else:
                    conn_parts.append(f"--{key}={value}")
        
        return " ".join(conn_parts)
    
    def check_database_exists(self, dbname: str) -> bool:
        """检查数据库是否存在"""
        try:
            conn_string = self.build_connection_string("postgres")
            # 设置环境变量和构建命令
            env = os.environ.copy()
            if 'password' in self.conn_params and self.conn_params['password']:
                env['PGPASSWORD'] = self.conn_params['password']
            
            cmd_parts = ['psql']
            if 'host' in self.conn_params:
                cmd_parts.append(f"--host={self.conn_params['host']}")
            if 'port' in self.conn_params:
                cmd_parts.append(f"--port={self.conn_params['port']}")
            if 'user' in self.conn_params:
                cmd_parts.append(f"--username={self.conn_params['user']}")
            cmd_parts.extend(['-tAc', f"SELECT 1 FROM pg_database WHERE datname = '{dbname}'"])
            cmd_parts.append('postgres')
            
            result = run_command(cmd_parts, capture_output=True, env=env)
            
            if result.returncode == 0 and result.stdout.strip() == "1":
                return True
            return False
        except Exception as e:
            logger.error(f"检查数据库存在性失败: {e}")
            return False
    
    def create_test_database(self) -> bool:
        """创建测试数据库"""
        try:
            if self.check_database_exists(self.test_db):
                print_info(f"数据库 {self.test_db} 已存在")
                if confirm("是否删除并重新创建测试数据库?", True):
                    self.drop_test_database()
                else:
                    return True
            
            print_info(f"创建测试数据库: {self.test_db}")
            
            # 设置环境变量
            env = os.environ.copy()
            if 'password' in self.conn_params and self.conn_params['password']:
                env['PGPASSWORD'] = self.conn_params['password']
            
            # 构建createdb命令
            cmd_parts = ['createdb']
            if 'host' in self.conn_params:
                cmd_parts.append(f"--host={self.conn_params['host']}")
            if 'port' in self.conn_params:
                cmd_parts.append(f"--port={self.conn_params['port']}")
            if 'user' in self.conn_params:
                cmd_parts.append(f"--username={self.conn_params['user']}")
            cmd_parts.append(self.test_db)
            
            result = run_command(cmd_parts, capture_output=True, env=env)
            
            if result.returncode == 0:
                print_success(f"测试数据库 {self.test_db} 创建成功")
                return True
            else:
                print_error(f"创建测试数据库失败: {result.stderr}")
                return False
                
        except Exception as e:
            print_error(f"创建测试数据库时出错: {e}")
            return False
    
    def drop_test_database(self) -> bool:
        """删除测试数据库"""
        try:
            if not self.check_database_exists(self.test_db):
                return True
                
            print_info(f"删除测试数据库: {self.test_db}")
            
            # 设置环境变量
            env = os.environ.copy()
            if 'password' in self.conn_params and self.conn_params['password']:
                env['PGPASSWORD'] = self.conn_params['password']
            
            # 构建dropdb命令
            cmd_parts = ['dropdb']
            if 'host' in self.conn_params:
                cmd_parts.append(f"--host={self.conn_params['host']}")
            if 'port' in self.conn_params:
                cmd_parts.append(f"--port={self.conn_params['port']}")
            if 'user' in self.conn_params:
                cmd_parts.append(f"--username={self.conn_params['user']}")
            cmd_parts.append(self.test_db)
            
            result = run_command(cmd_parts, capture_output=True, env=env)
            
            if result.returncode == 0:
                print_success(f"测试数据库 {self.test_db} 删除成功")
                return True
            else:
                print_error(f"删除测试数据库失败: {result.stderr}")
                return False
                
        except Exception as e:
            print_error(f"删除测试数据库时出错: {e}")
            return False
    
    def initialize_pgbench(self, scale_factor: int) -> bool:
        """初始化pgbench测试数据"""
        try:
            print_info(f"正在初始化pgbench测试数据，规模因子: {scale_factor}")
            
            # 设置环境变量
            env = os.environ.copy()
            if 'password' in self.conn_params and self.conn_params['password']:
                env['PGPASSWORD'] = self.conn_params['password']
            
            # 构建pgbench初始化命令 - 修复参数格式
            cmd_parts = ['pgbench', '-i']
            
            # 添加连接参数
            if 'host' in self.conn_params:
                cmd_parts.extend(['-h', self.conn_params['host']])
            if 'port' in self.conn_params:
                cmd_parts.extend(['-p', str(self.conn_params['port'])])
            if 'user' in self.conn_params:
                cmd_parts.extend(['-U', self.conn_params['user']])
            
            # 添加初始化参数
            cmd_parts.extend(['-s', str(scale_factor)])
            
            # 添加数据库名
            cmd_parts.append(self.test_db)
            
            print_info(f"初始化命令: {' '.join(cmd_parts)}")
            
            # 使用进度条显示初始化进度
            with ProgressManager(description="初始化测试数据") as progress:
                result = run_command(cmd_parts, capture_output=True, env=env)
                progress.update(100)  # 完成进度
            
            print_info(f"初始化返回码: {result.returncode}")
            if result.stderr:
                print_info(f"初始化错误: {result.stderr}")
            
            if result.returncode == 0:
                print_success("pgbench测试数据初始化成功")
                return True
            else:
                print_error(f"pgbench初始化失败: {result.stderr}")
                return False
                
        except Exception as e:
            print_error(f"初始化pgbench数据时出错: {e}")
            return False
    
    def run_pgbench_test(self, scale_factor: int, clients: int, threads: int, duration: int, 
                        test_type: str = "read-write") -> Optional[PgbenchTestResult]:
        """运行pgbench测试"""
        try:
            print_info(f"开始pgbench {test_type} 测试...")
            print_info(f"参数: 规模因子={scale_factor}, 客户端={clients}, 线程={threads}, 时长={duration}秒")
            
            # 设置环境变量
            env = os.environ.copy()
            if 'password' in self.conn_params and self.conn_params['password']:
                env['PGPASSWORD'] = self.conn_params['password']
            
            # 先测试简单的连接
            print_info("测试数据库连接...")
            test_cmd = ['psql']
            if 'host' in self.conn_params:
                test_cmd.extend(['-h', self.conn_params['host']])
            if 'port' in self.conn_params:
                test_cmd.extend(['-p', str(self.conn_params['port'])])
            if 'user' in self.conn_params:
                test_cmd.extend(['-U', self.conn_params['user']])
            test_cmd.extend(['-c', 'SELECT 1', self.test_db])
            
            test_result = run_command(test_cmd, capture_output=True, env=env)
            print_info(f"连接测试返回码: {test_result.returncode}")
            if test_result.stderr:
                print_info(f"连接测试错误: {test_result.stderr}")
            
            if test_result.returncode != 0:
                print_error("数据库连接失败，请检查连接参数和权限")
                return None
            
            # 构建pgbench测试命令 - 使用更简单的参数
            cmd_parts = ['pgbench']
            
            # 添加连接参数
            if 'host' in self.conn_params:
                cmd_parts.extend(['-h', self.conn_params['host']])
            if 'port' in self.conn_params:
                cmd_parts.extend(['-p', str(self.conn_params['port'])])
            if 'user' in self.conn_params:
                cmd_parts.extend(['-U', self.conn_params['user']])
            
            # 添加测试参数 - 先用简单参数测试
            cmd_parts.extend(['-c', str(clients)])
            cmd_parts.extend(['-j', str(threads)])
            cmd_parts.extend(['-T', str(duration)])
            
            # 对于测试类型，使用简单的内置类型
            builtin_types = ['simple-update', 'select-only']
            if test_type in builtin_types:
                cmd_parts.extend(['-b', test_type])
            
            # 添加数据库名
            cmd_parts.append(self.test_db)
            
            print_info(f"执行命令: {' '.join(cmd_parts)}")
            print_info("正在执行性能测试，请稍候...")
            
            # 执行测试并捕获输出
            start_time = time.time()
            result = run_command(cmd_parts, capture_output=True, env=env)
            end_time = time.time()
            
            print_info(f"命令执行时间: {end_time - start_time:.2f}秒")
            print_info(f"返回码: {result.returncode}")
            
            if result.stdout:
                print_info(f"标准输出: {result.stdout}")
            if result.stderr:
                print_info(f"标准错误: {result.stderr}")
            
            if result.returncode == 0:
                # 解析pgbench输出
                test_result = self.parse_pgbench_output(result.stdout, test_type, scale_factor, clients, threads, duration)
                if test_result:
                    self.test_results.append(test_result)
                    print_success(f"pgbench {test_type} 测试完成")
                    return test_result
                else:
                    print_error("解析pgbench输出失败")
                    return None
            else:
                print_error(f"pgbench测试失败 (返回码: {result.returncode})")
                print_error("可能的原因:")
                print_error("1. pgbench版本不支持-b参数")
                print_error("2. 数据库连接权限不足")
                print_error("3. 测试数据未正确初始化")
                print_error("4. 客户端数量过多")
                
                # 尝试不带-b参数的简单测试
                print_info("尝试不带-b参数的简单测试...")
                simple_cmd = [part for part in cmd_parts if part != '-b' and part != test_type]
                simple_result = run_command(simple_cmd, capture_output=True, env=env)
                
                print_info(f"简单测试返回码: {simple_result.returncode}")
                if simple_result.stdout:
                    print_info(f"简单测试输出: {simple_result.stdout}")
                if simple_result.stderr:
                    print_info(f"简单测试错误: {simple_result.stderr}")
                
                if simple_result.returncode == 0:
                    print_success("简单测试成功！使用默认测试类型")
                    test_result = self.parse_pgbench_output(simple_result.stdout, "default", scale_factor, clients, threads, duration)
                    if test_result:
                        self.test_results.append(test_result)
                        return test_result
                
                return None
                
        except Exception as e:
            print_error(f"运行pgbench测试时出错: {e}")
            return None
    
    def parse_pgbench_output(self, output: str, test_type: str, scale_factor: int, 
                           clients: int, threads: int, duration: int) -> Optional[PgbenchTestResult]:
        """解析pgbench输出"""
        try:
            print_info("解析pgbench输出...")
            print_info(f"原始输出: {output}")
            
            # 使用正则表达式提取关键指标 - 根据实际输出格式更新
            patterns = {
                'tps_excluding': r'tps\s*=\s*([\d.]+)\s*\(without initial connection time\)',
                'tps_including': r'tps\s*=\s*([\d.]+)\s*\(including connections establishing\)',
                'latency_avg': r'latency average\s*=\s*([\d.]+)\s*ms',
                'latency_stddev': r'latency stddev\s*=\s*([\d.]+)\s*ms',
                'transactions': r'number of transactions actually processed:\s*(\d+)',
                'failed_transactions': r'number of failed transactions:\s*(\d+)\s*\(([\d.]+)%\)'
            }
            
            values = {}
            for key, pattern in patterns.items():
                match = re.search(pattern, output)
                if match:
                    values[key] = float(match.group(1))
                    logger.info(f"成功解析 {key}: {values[key]}")
                else:
                    logger.warning(f"无法解析pgbench输出中的: {key}")
                    values[key] = 0.0
            
            # 如果没有找到TPS，尝试计算
            if values['tps_excluding'] == 0 and values['transactions'] > 0:
                values['tps_excluding'] = values['transactions'] / duration
                logger.info(f"计算得到TPS: {values['tps_excluding']}")
            
            # 如果没有找到包含连接的TPS，使用排除连接的TPS
            if values['tps_including'] == 0:
                values['tps_including'] = values['tps_excluding']
            
            # 如果没有找到延迟标准差，使用平均值的一半作为估算
            if values['latency_stddev'] == 0 and values['latency_avg'] > 0:
                values['latency_stddev'] = values['latency_avg'] * 0.3
                logger.info(f"估算延迟标准差: {values['latency_stddev']}")
            
            print_info("解析结果:")
            for key, value in values.items():
                print_info(f"  {key}: {value}")
            
            return PgbenchTestResult(
                test_type=test_type,
                scale_factor=scale_factor,
                clients=clients,
                threads=threads,
                duration=duration,
                tps_including_connections=values['tps_including'],
                tps_excluding_connections=values['tps_excluding'],
                latency_stddev=values['latency_stddev'],
                latency_avg=values['latency_avg'],
                success=True
            )
            
        except Exception as e:
            logger.error(f"解析pgbench输出失败: {e}")
            logger.error(f"输出内容: {output}")
            return None
    
    def display_test_result(self, result: PgbenchTestResult):
        """显示单个测试结果"""
        table = Table(title=f"pgbench {result.test_type.upper()} 测试结果", box=ROUNDED)
        table.add_column("指标", style="cyan")
        table.add_column("值", style="green", justify="right")
        table.add_column("单位", style="yellow")
        
        table.add_row("测试类型", result.test_type, "")
        table.add_row("规模因子", str(result.scale_factor), "")
        table.add_row("客户端数", str(result.clients), "")
        table.add_row("线程数", str(result.threads), "")
        table.add_row("测试时长", str(result.duration), "秒")
        table.add_row("TPS (包含连接)", f"{result.tps_including_connections:.2f}", "事务/秒")
        table.add_row("TPS (排除连接)", f"{result.tps_excluding_connections:.2f}", "事务/秒")
        table.add_row("平均延迟", f"{result.latency_avg:.2f}", "毫秒")
        table.add_row("延迟标准差", f"{result.latency_stddev:.2f}", "毫秒")
        
        console.print(table)
    
    def display_all_results(self):
        """显示所有测试结果"""
        if not self.test_results:
            print_warning("没有测试结果可显示")
            return
        
        # 创建比较表格
        table = Table(title="pgbench 性能测试结果汇总", box=ROUNDED)
        table.add_column("测试类型", style="cyan")
        table.add_column("规模因子", style="green", justify="right")
        table.add_column("客户端", style="green", justify="right")
        table.add_column("线程", style="green", justify="right")
        table.add_column("TPS", style="yellow", justify="right")
        table.add_column("平均延迟", style="yellow", justify="right")
        
        for result in self.test_results:
            table.add_row(
                result.test_type,
                str(result.scale_factor),
                str(result.clients),
                str(result.threads),
                f"{result.tps_excluding_connections:.2f}",
                f"{result.latency_avg:.2f}ms"
            )
        
        console.print(table)
    
    def get_test_parameters(self) -> Tuple[int, int, int, int]:
        """获取测试参数"""
        print_info("配置pgbench测试参数")
        
        # 获取规模因子
        scale_factor = self._get_positive_int_input("请输入规模因子 (1=1MB, 10=10MB, 100=100MB)", 10)
        
        # 获取客户端数量
        clients = self._get_positive_int_input("请输入客户端数量", 4)
        
        # 获取线程数量
        threads = self._get_positive_int_input("请输入线程数量", 2)
        
        # 获取测试时长
        duration = self._get_positive_int_input("请输入测试时长(秒)", 60)
        
        # 验证参数合理性
        if threads > clients:
            print_warning("线程数不应大于客户端数，自动调整为客户端数")
            threads = clients
        
        return scale_factor, clients, threads, duration
    
    def get_test_type(self) -> str:
        """获取测试类型"""
        print_info("选择测试类型")
        print("1. default (默认测试，TPC-B-like)")
        print("2. simple-update (简单更新测试)")
        print("3. select-only (只读测试)")
        
        choice = get_input("请选择测试类型", "1")
        
        type_map = {
            "1": "default",
            "2": "simple-update", 
            "3": "select-only"
        }
        
        return type_map.get(choice, "default")
    
    def _get_positive_int_input(self, prompt: str, default: int) -> int:
        """获取正整数输入"""
        while True:
            value = get_input(prompt, str(default))
            try:
                int_value = int(value)
                if int_value > 0:
                    return int_value
                print_error("请输入大于0的整数")
            except ValueError:
                print_error("请输入有效的整数")

class PgbenchManager:
    """pgbench性能测试管理类"""
    
    def __init__(self):
        """初始化pgbench管理器"""
        self.tester = PgbenchTester()
    
    def start_pgbench_test(self):
        """启动pgbench性能测试"""
        print_info("准备启动pgbench性能测试...")
        
        # 检查pgbench是否可用
        if not self.tester.check_pgbench_availability():
            print_error("pgbench不可用，请确保PostgreSQL已正确安装")
            return
        
        # 获取数据库连接参数
        if not self.tester.get_connection_params():
            print_error("获取数据库连接参数失败")
            return
        
        # 创建测试数据库
        if not self.tester.create_test_database():
            print_error("创建测试数据库失败")
            return
        
        try:
            # 获取测试参数
            scale_factor, clients, threads, duration = self.tester.get_test_parameters()
            
            # 初始化测试数据
            if not self.tester.initialize_pgbench(scale_factor):
                print_error("初始化测试数据失败")
                return
            
            # 执行多个测试
            test_types = []
            if confirm("是否执行默认测试?", True):
                test_types.append("default")
            if confirm("是否执行简单更新测试?", False):
                test_types.append("simple-update")
            if confirm("是否执行只读测试?", False):
                test_types.append("select-only")
            
            if not test_types:
                print_warning("未选择任何测试类型，默认执行默认测试")
                test_types = ["default"]
            
            # 执行测试
            for test_type in test_types:
                result = self.tester.run_pgbench_test(scale_factor, clients, threads, duration, test_type)
                if result:
                    self.tester.display_test_result(result)
                
                # 测试间隔
                if test_type != test_types[-1]:  # 不是最后一个测试
                    if confirm("是否继续下一个测试?", True):
                        print_info("等待5秒后开始下一个测试...")
                        time.sleep(5)
                    else:
                        break
            
            # 显示所有结果
            if len(self.tester.test_results) > 1:
                self.tester.display_all_results()
            
        except KeyboardInterrupt:
            print_warning("\n测试被用户中断")
        except Exception as e:
            print_error(f"测试过程中出错: {e}")
            logger.exception("pgbench测试出错")
        finally:
            # 询问是否清理测试数据
            if confirm("是否删除测试数据库?", True):
                self.tester.drop_test_database()
            else:
                print_warning(f"测试数据库 {self.tester.test_db} 未删除")
                print_info(f"可以手动执行: dropdb {self.tester.test_db}")

def main():
    """pgbench测试模块主函数"""
    manager = PgbenchManager()
    manager.start_pgbench_test()

if __name__ == "__main__":
    main()